import torch
import torch.nn as nn
import torch.optim as optim
import random
import re

# 设置随机种子以保证结果可复现
torch.manual_seed(42)
random.seed(42)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# 准备数据
text = '''亲爱的朋友们，我所站立的地方，就是中国。

　　如果你侧耳倾听，你是否听到了那田间蛙鸣和大漠中的驼铃？

　　俯瞰中国乡村的上空，田间依旧的蛙鸣见证了美丽乡村的嬗变。中华有诗云：“稻花香里说丰年，听取蛙声一片。”从关注“三农”到精准扶贫，时代的巨轮飞速向前时，乡村不曾被抛下。它以非凡的活力实现了自身的蜕变，呈现出一幅山水美、人情美的淳朴画卷。

　　从微观层面上，无数致力于美丽乡村建设的人们成为时代的先锋。耶鲁大学村官秦玥飞带领乡村率先走上信息化的轨道；三十年支教女教师以芳兰振蕙叶之姿，使“支月英”这个名字镌刻在两代人的记忆之中。鲁迅曾说：“无穷的远方，无尽的人们，都与我有关。”正是怀着对乡村大地最深沉的眷恋和责任，无数人与曾经的这方土地有了再难分割的羁绊。从宏观层面上，国家的发展从未遗忘农村，从顶层设计再到基层探索，湘西十八洞村的扶贫经验至今还在传播，三权分置体现的对农民的关怀也仍是人们感念的对象。

　　大风起于青萍之末，知止而后有定，激荡之中而后有静。

　　放眼宏观的国际，新丝路上又响起的阵阵驼铃，传来人类命运共同体的集体乡愁。西方有谚云“独木不成林”，非洲朋友也说“水涨荷花高”，一带一路的“传奇”构想是古老中国全球担当的表征，它不是以中国霸权睥睨一切，而是以开放的怀抱包容世界。

　　当普惠互利共赢成为全球治理的共同呼声，中国成为最勇敢的弄潮儿。当“逆全球化”的言论甚嚣尘上，当地球村的乐观之幕徐徐降下，有人发出“世界怎么了？我们怎么办？”的忧患之问。一带一路成为全球对于中国治理方案的共同期待。中国人的文化自觉与文化自信凝聚成责任担当，古老的丝绸之路是养分之源，全新的发展理念又为其注入了不竭的活力。

　　当蛙声伴着田间的风传遍广阔的美丽乡村，当驼铃载着国际的梦奏响一带一路之歌，你会看见古老中国日益崛起的身姿，它化为一声厚重的龙吟，泅渡了历史的重洋，激荡在七大洲的上空。

　　我所站立的地方，就是我的中国，我是什么，中国便是什么。我怎么样，中国便怎么样。

　　今日少年人，当蛙声伴着驼铃，我们向前迈步的脚步声应当更加铿锵有力才是。

　　世界青年人，也请让我们伸出手，握一握，共同穿越巨变的中国。'''

text = re.sub(r'[^\u4e00-\ufa95，。！？]', '', text)

chars = sorted(list(set(text)))
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}


# 将文本转换为索引
def text_to_indices(text, char_to_idx):
    return [char_to_idx[ch] for ch in text]


# 准备训练数据
seq_length = 20
data = text_to_indices(text, char_to_idx)
X = []
Y = []

for i in range(len(data) - seq_length):
    X.append(data[i:i + seq_length])
    Y.append(data[i + seq_length])

X = torch.tensor(X, dtype=torch.long)
Y = torch.tensor(Y, dtype=torch.long)

from torch.utils.data import Dataset, DataLoader


class CharDataset(Dataset):
    def __init__(self, X, Y):
        self.X = X
        self.Y = Y

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.Y[idx]


# 创建数据集和数据加载器
batch_size = 8  # 设置你想要的批量大小
dataset = CharDataset(X, Y)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)


# 定义 LSTM 模型
class CharLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers):
        super(CharLSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, device=device)
        self.fc = nn.Linear(hidden_size, output_size, device=device)

    def forward(self, x, hidden):
        out, hidden = self.lstm(x, hidden)
        out = self.fc(out[:, -1, :])  # 取序列的最后一个时间步
        return out, hidden

    def init_hidden(self, batch_size):
        h0 = torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device)
        c0 = torch.zeros(self.num_layers, batch_size, self.hidden_size).to(device)
        return h0, c0


# 模型参数
input_size = len(chars)  # 输入大小（字符种类数）
hidden_size = 128
output_size = len(chars)  # 输出大小（字符种类数）s
num_layers = 1

model = CharLSTM(input_size, hidden_size, output_size, num_layers).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)


# 训练模型
def train(model, X, Y, epochs=100):
    model.train()

    for epoch in range(epochs):
        total_loss = 0

        for batch_x, batch_y in dataloader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)

            batch_size = batch_x.size(0)
            hidden = model.init_hidden(batch_size=batch_size)

            # One-hot 编码
            batch_x = batch_x.unsqueeze(2)  # [batch_size, seq_length, 1]
            batch_x = torch.zeros(batch_x.size(0), batch_x.size(1), input_size).to(device).scatter_(2, batch_x, 1)

            # 前向传播
            output, hidden = model(batch_x.float(), hidden)
            loss = criterion(output, batch_y)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        if (epoch + 1) % 10 == 0:
            print(f"Epoch [{epoch + 1}/{epochs}], Loss: {total_loss / len(dataloader):.4f}")


# 训练模型
train(model, X, Y, epochs=100)


# 生成文本
def generate_text(model, start_str, length=100):
    model.eval()
    hidden = model.init_hidden(batch_size=1)
    chars = list(start_str)

    for p in range(len(chars) - 1):
        x = torch.tensor([[char_to_idx[chars[p]]]]).to(device).unsqueeze(2)
        x = torch.zeros(1, 1, input_size).to(device).scatter_(2, x, 1)
        _, hidden = model(x, hidden)

    pred = chars[-1]

    for p in range(length):
        x = torch.tensor([[char_to_idx[pred]]]).to(device).unsqueeze(2)
        x = torch.zeros(1, 1, input_size).to(device).scatter_(2, x, 1)
        output, hidden = model(x, hidden)

        output_dist = output.data.view(-1).exp()
        top_i = torch.multinomial(output_dist, 1)[0]
        pred = idx_to_char[top_i.item()]
        chars.append(pred)

    return ''.join(chars)


# 生成文本
print("Generated text:")
print(generate_text(model, start_str="朋友", length=200))

torch.save(model.state_dict(), "pytorch-lstm-word.pth")  # 保存模型参数
